package lsieun.visitor;

import lsieun.cst.Constant;
import lsieun.utils.ContainerUtils;
import lsieun.utils.NameUtils;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;

public class PatchTransformVisitor extends ClassVisitor {
    public PatchTransformVisitor(ClassVisitor classVisitor) {
        super(Constant.API_VERSION, classVisitor);
    }

    @Override
    public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
        MethodVisitor mv = super.visitMethod(access, name, descriptor, signature, exceptions);
        if (mv != null &&
                name.equals("transform") &&
                descriptor.equals("(Ljava/lang/ClassLoader;Ljava/lang/String;Ljava/lang/Class;Ljava/security/ProtectionDomain;[B)[B")) {
            mv = new PatchTransformAdapter(mv);
        }
        return mv;
    }

    private static class PatchTransformAdapter extends MethodVisitor implements Opcodes {
        private final String containerName = NameUtils.toInternalName(ContainerUtils.class.getName());

        public PatchTransformAdapter(MethodVisitor methodVisitor) {
            super(Constant.API_VERSION, methodVisitor);
        }

        @Override
        public void visitCode() {
            if (mv != null) {
                mv.visitVarInsn(ALOAD, 2);
                mv.visitVarInsn(ALOAD, 5);
                mv.visitMethodInsn(INVOKESTATIC, containerName, "addClass", "(Ljava/lang/String;[B)V", false);
            }

            super.visitCode();
        }

        @Override
        public void visitInsn(int opcode) {
            if (mv != null) {
                if (opcode == ARETURN) {
                    mv.visitInsn(DUP);
                    mv.visitVarInsn(ALOAD, 2);
                    mv.visitInsn(SWAP);
                    mv.visitMethodInsn(INVOKESTATIC, containerName, "compareClass", "(Ljava/lang/String;[B)V", false);
                }
            }
            super.visitInsn(opcode);
        }
    }
}
